#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""

"""

import copy

class SuffixTreeNode(dict):
    def __init__(self, s='', start_idx=-1, parent=None):
        self.s = s
        # self.children = {}  # {s : SuffixTreeNode}
        self.leaf = start_idx
        self.parent = parent

    def add_child(self, s, idx):
        for child_s, child_node in self.items():
           if child_node.merge(s, idx):
                return child_s

        self[s] = SuffixTreeNode(s, idx, parent=self)

    def merge(self, target_s, start_idx):
        # target_s : 待插入的字符串.

        # ret_tuple == (最终匹配到了哪个节点匹配不下去了, 匹配到节点的哪个索引, 字符串剩余内容)
        match_node, match_node_mat_idx, remain_target_s = self._match(target_s)

        if match_node_mat_idx == -1:
            # 当前节点没有匹配到target_s, 所以target_s, 不能插入到当前节点
            return False
        else:
            # 当前节点匹配到target_s一部分, 现在插入
            # match_node 要分裂
            match_node_left  = match_node.s[0:match_node_mat_idx+1]
            match_node_right = match_node.s[match_node_mat_idx+1:match_node.s_len]

            # split
            # print('left %s, right %s' % (match_node_left, match_node_right))
            if match_node_right == '':
                # match_node 的leaf性质不变, 既可以作为leaf, 也可以作为parent
                # print('===================  =====================')
                if remain_target_s:
                    match_node[remain_target_s] = SuffixTreeNode(remain_target_s, start_idx, parent=match_node)
                else:
                    match_node.leaf = start_idx
            else:
                # split
                # assert remain_target_s
                # print('not full')
                new_child = SuffixTreeNode(s=match_node_right, parent=match_node)
                new_child = copy.deepcopy(match_node)

                new_child.leaf = match_node.leaf
                new_child.s = match_node_right
                new_child.parent = match_node
                # new_child.leaf = match_node_right

                # seg new s and children
                match_node.clear()

                # match_node.s and parent's link
                del match_node.parent[match_node.s]
                match_node.s = match_node_left
                match_node.parent[match_node_left] = match_node

                # print(id (match_node))
                match_node.leaf = -1

                # add child
                match_node[match_node_right] = new_child

                if remain_target_s:
                    match_node[remain_target_s] = SuffixTreeNode(remain_target_s, start_idx, match_node)
                else:
                    match_node.leaf = start_idx

            return True

    def _match(self, target_s):
        # match_last_idx 是最后一个匹配到的索引
        # return match_node, match_node_match_idx, remain_target_s
        match_last_idx = self.str_match_len(self.s, target_s)

        assert match_last_idx +1 <= self.s_len

        remain_target_s = target_s[match_last_idx+1:]
        if match_last_idx + 1 < self.s_len or remain_target_s == '':
            # 没有匹配完当前字符
            return self, match_last_idx, remain_target_s

        # 当前节点self.s全部内容都和target_s匹配了, 匹配子节点
        assert match_last_idx + 1 == self.s_len
        # 匹配完当前节点字符串, 去匹配子节点
        for child in self.values():
            child_node, child_match_idx, child_remain_target_s = child._match(remain_target_s)
            if child_match_idx != -1:
                return child_node, child_match_idx, child_remain_target_s

        # 如果子节点都没有匹配到, 返回当前节点匹配结果
        return self, match_last_idx, remain_target_s

    def match(self, pattern):
        for child in self.values():
            match_node, match_node_idx, remain = child._match(pattern)
            if match_node_idx != -1:
                # print('match_s %s' % match_node.s)
                # print('match_node_idx %s' % match_node_idx)
                # print('remain %s' % remain)
                return remain == ''

        return False

    def match_cnt(self, pattern):
        for child in self.values():
            match_node, match_node_idx, remain = child._match(pattern)
            if match_node_idx != -1:
                # print('match_s %s' % match_node.s)
                # print('match_node_idx %s' % match_node_idx)
                # print('remain %s' % remain)
                if remain != '':
                    return 0
                else:
                    return match_node.leaf_cnt()

        return 0

    def leaf_cnt(self):
        child_cnt = 0

        for child in self.values():
            child_cnt += child.leaf_cnt()

        if self.leaf != -1:
            return child_cnt + 1
        else:
            return child_cnt


    @property
    def s_len(self):
        return len(self.s)

    def str_match_len(self, s1, s2):
        # 如果不配, 返回-1
        # print('%s %s' % (s1, s2))
        len_1 = len(s1)
        len_2 = len(s2)

        little_len = min(len_1, len_2)
        i = 0

        for i in range(little_len):
            if s1[i] != s2[i]:
                return i-1
        return i


class SuffixTree:
    '''
    压缩suffix tree
    本质还是一颗trie树
    '''

    def __init__(self, s):
        self.root = SuffixTreeNode()

        # for i in range(len(s)):
        for i in range(len(s)-1, -1, -1):
            # print('_insert %s , idx  %s ' % (s[i:], i))
            self._insert(s[i:], i)

    def _insert(self, s, idx):
        self.root.add_child(s, idx)

    def match(self, s):
        return self.root.match(s)

    def match_cnt(self, s):
        return self.root.match_cnt(s)

    def travel(self):
        from pprint import pprint

        pprint(self.root)
        self._dfs_travel(self.root, '')

    def _dfs_travel(self, node, path):
        if node is None:
            return

        if node.leaf != -1:
            print('idx %s, s %s' % (node.leaf, path+node.s))

        for child in node.values():
            self._dfs_travel(child, path + node.s)



def main():
    print("start main")
    # s = SuffixTree('abcabc')
    # s = SuffixTree('abc')

    # 测试用例来自: https://www.youtube.com/watch?v=VA9m_l6LpwI
    s = SuffixTree('CAGTCAGG')

    # print(list(s.root.keys() ))
    c = s.root
    s.travel()

    ret = s.match('C')
    print('C %s' % ret)

    ret = s.match('D')
    print('D %s' % ret)

    ret = s.match_cnt('G')
    print('G match_cnt %s' % ret)

    # print('===================  =====================')
    # ret = s.match('GTCAGG')
    # print('GTCAGGG %s' % ret)

    # print(c['a']['a']['a'])

    # print(c)
    # print(id(c['a']))
    # print(c['a'].s)
    # print(c['a'].leaf)
    # print(c['a']['a'].leaf)
    # print(c['a']['a']['a'].leaf)
    # print(c['a']['a']['a']['a'].leaf)

if __name__ == "__main__":
    main()
